Conversation
The neighborhood gufuncs are built with ``cache=True``, but the cache never hit: every process recompiled all nine kernels and appended a fresh set of entries to the index, which grew without bound (over a thousand files in a working checkout). The cause is that ``_make_kernel`` closed each kernel over its reducer. Numba keys a cached function on a hash of its closure, and a ``Dispatcher`` serializes with a ``uuid4`` regenerated in every process, so the key differed in every process. Numba reports no error for this -- it writes the cache and silently misses it. All nine kernels also shared one cache file, since the closure made them one function as far as numba's locator was concerned. Give each reduction a module-level kernel body instead, and apply ``guvectorize`` to it on first call. A module-level body has no closure, so its key is stable, and it reaches its reducer as a global, which numba resolves at compile time and leaves out of the key. Each body also gets its own cache file. Kernels are now compiled once per machine rather than once per process. Compilation stays lazy. ``guvectorize`` compiles at decoration time when given explicit signatures, so decorating at module scope would rebuild every kernel during ``import uxarray`` and start numba's threading layer, leaving a thread pool that makes forks unsafe -- the regression ``test_no_numba_kernels_built_on_import`` guards against. Replace ``functools.cache`` with a double-checked lock while here. It does not hold a lock across the call it memoizes, so under ``dask="parallelized"`` every worker thread that reached a kernel before the first build finished started its own full compilation, serialized behind numba's global compiler lock. A 12-chunk reduction compiled the same kernel 12 times; it now compiles once. The bodies are spelled out rather than generated because the deduplication one would reach for -- a single shared body taking the reducer as an argument -- makes the reducer a dynamic global, which numba refuses to cache at all. The gather is shared through ``_widest``/``_gather`` instead, leaving only the reducer name different between bodies. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Replaces the ``_lazy`` closure with a ``_LazyKernel`` descriptor that holds the compiled gufunc on itself, so each kernel is still built on first use but the memoization lives at class level rather than in a per-kernel closure. Two things fall out. The ``staticmethod`` wrapping goes away: a descriptor hands back the gufunc itself, so ``self`` is never bound as the kernel's first argument, and the comment explaining that wart goes with it. And compilation now happens where the attribute is resolved -- ``_apply_kernel`` reads it on the calling thread -- so a dask-backed reduction builds its kernel while the graph is being assembled rather than inside a task. The array stays lazy; only the compile moves, and it moves out of the parallel region, which makes the worker-thread race structurally impossible rather than merely locked against. ``functools.cached_property`` would not do. It caches per instance, and ``Grid.neighborhood()`` returns a new ``Neighborhood`` on every call, so it ran ``guvectorize`` once per neighborhood: fifty fresh instances measured fifty compilations, against zero for the descriptor. It also holds no lock, CPython having removed it in 3.12, so twelve threads racing one instance measured twelve compilations, against one. The kernel bodies must still be module-level functions with no closure. That is what keeps numba's cache key stable across processes, and it is independent of where the memoization lives -- the memoizer never enters the key. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
Sevans711
left a comment
There was a problem hiding this comment.
The fix looks clean, even though it might look like there are some "object-oriented shenanigans" happening here with making _LazyKernel class a descriptor (i.e., it has a __get__ method), it makes sense to me and the comments communicate the intent clearly, so it shouldn't be too difficult to maintain in the future if needed. I think this is a good way to proceed.
Noticed the small slowdowns (~20%) in some neighborhood benchmarks, but I think that is acceptable if this PR is solving the race condition and/or lack of caching (which slows down real workflows the first time it hits but isn't necessarily appearing in the benchmark suite).
The only issue is, I still have the crash from the race condition (#1757)! Copy-pasting those lines into terminal while on this PR's branch, it still crashes locally for me within a few runs.
If you want to merge this as closing #1769 and related to 1757 but not closing 1757, I would be okay to approve. But for closing 1757 something else needs to be fixed. I'm not sure what is missing (race conditions are tricky to debug…) but I wonder if it might be as simple as moving the threading.Lock() mechanism (not sure where it would need to move to, though…)?
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Sevans711
left a comment
There was a problem hiding this comment.
Looks like the race condition is now fixed as well! Running the code from #1757, I now see no crashes after 20 runs, whereas before it was usually crashing within a few runs at most, and never took more than 10 runs before the race condition caused a crash.
Skimmed through the code changes since last time I viewed and they look reasonable.
My only concern now is that some of the 480km Neighborhood benchmarks show slowdowns. Though, the 120km ones show speedups…? This makes me wonder if maybe the ways this PR affects when numba compilation happens here is causing the ASV benchmarking suite to not warm up properly in this case. Are these slowdowns real, or just because they are including the numba compilation times? Aside from that, I would be happy to approve!
|
@Sevans711 Good catch, there was some numba compilation getting into the timings on smaller grids. Otherwise, there is a small overall speedup on the neighborhood reductions. Some of the other benchmark differences are probably noise (like the |
Sevans711
left a comment
There was a problem hiding this comment.
Approving, now that the benchmarks look better!
Closes #1757, #1769
Overview
The PR solves two issues with the current version of the neighborhood filters kernels:
njit(cache=True). This means that in the current implementation of kernel caching, every call to a kernel gets a new UUID, so any attempt to cache this way will basically be a no-op.In a fun twist, both are known unresolved issues in numba:
numba/numba#6264
numba/numba#9288
PR Checklist
General
Testing & Benchmarking
Documentation and Examples
docs/api.rst; internal (private) function names start with an underscore (_)AI Disclosure
AI Usage: Claude Opus 5